#!/usr/bin/python -i

from ceorModule import *

### Gas2D History
class Gas2DHistory:
  def __init__(self,total_time, each_time, gas):
    if isinstance(obj, MyClass):
      self.gas = gas
    else:
      print("bad boy, very bad voy!, it Must be a Gas2D class!")
      sys.exit()
    self.ttime = range(total_time)
    self.timeDim = total_time #/ttime_print + 1
    self.xt  = np.zeros([gas.N, timeDim])
    self.yt  = np.zeros([gas.N, timeDim])
    self.vxt = np.zeros([gas.N, timeDim])
    self.vyt = np.zeros([gas.N, timeDim])
    self.fxt = np.zeros([gas.N, timeDim])
    self.fyt = np.zeros([gas.N, timeDim])
    self.Ekint = np.zeros(timeDim)
    self.Epott = np.zeros(timeDim)
    self.Etott = np.zeros(timeDim)
    
  def savetime(self, gas, itime):
    self.Ekint[i_print] = gas.Ekin
    self.Epott[i_print] = gas.Epot
    #save particle trajectories
    for j in range(gas.N):
      self.xt[j][i_print]  = gas.particle[j].x_new
      self.yt[j][i_print]  = gas.particle[j].y_new
      self.vxt[j][i_print] = gas.particle[j].vx_new
      self.vyt[j][i_print] = gas.particle[j].vy_new
      self.fxt[j][i_print] = gas.particle[j].fx_new
      self.fyt[j][i_print] = gas.particle[j].fy_new    
    

######################################
################ MAIN ################
## Is a 2-D box with lenght xL yL
if __name__ == "__main__":
  # And God says...
  N  = 10
  KT_eq = 0.1
  eps = 1.0
  sigma = 1.0
  delta_r_ini = 2.5*sigma
  dt = 0.0001
  gas = Gas2D(10, 5, N, KT_eq, delta_r_ini)
  prefix="N_0010_Part_0.0001"
  ext = ".png"
  for i in range(gas.N):
    gas.particle[i].vx_new = 0.0
    gas.particle[i].vy_new = 0.0
  
  #--- Plot initial positions
  gas.arrayPos()
  gas.arrayVel()
  print gas.X
  print gas.Y
  
  plt.figure(1)
  plt.grid()
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  plt.scatter(gas.X, gas.Y, marker='s')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  # to be continued...  
  #outputfilename=prefix+"Fig01-XY_ini_end"+ext
  #plt.savefig(outputfilename)
  
  # 1. Calculamos F(t)
  gas.calculateForces(eps, sigma)
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_ini = ', gas.Ekin, 'Epot_ini = ', gas.Epot

  # ------ Begin SIMULATION ------
  total_time = 100000 #4900
  ttime_print = 1
  ttime = range(total_time)
  timeDim = total_time #/ttime_print + 1
  xt  = np.zeros([gas.N, timeDim])
  yt  = np.zeros([gas.N, timeDim])
  vxt = np.zeros([gas.N, timeDim])
  vyt = np.zeros([gas.N, timeDim])
  fxt = np.zeros([gas.N, timeDim])
  fyt = np.zeros([gas.N, timeDim])
  Ekint = np.zeros(timeDim)
  Epott = np.zeros(timeDim)
  Etott = np.zeros(timeDim)
  #######str(n).zfill(6)

  for i in ttime:
    gas.save_old() # fx_old=fx_new
    gas.stepPosition(dt) # input old variables = > output x_new
    gas.inBox()
    gas.calculateForces(eps, sigma)
    gas.stepVelocity(dt)
    # Save every each ttime_print
    #if i % ttime_print == 0 :
    i_print = i #/ttime_print
    gas.calculateEkin()
    # gasHist.save(i_print, gas)

  # ------ End SIMULATION ------
  ttime = dt*np.array(ttime)
  Etott = Ekint+Epott
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_end = ', gas.Ekin, 'Epot_end = ', gas.Epot
  print '-- deltas --'*3
  print eps, dt, total_time, Etott.mean(), Etott.max() - Etott.min(),Etott.max(), Etott.min()
  
  
  # Save in file
  outdatafile=prefix+"Energy.dat"
  np.savetxt(outdatafile,zip(ttime,Epott, Ekint, Etott))
  
  gas.arrayPos()
  gas.arrayVel()
  plt.rc('text', usetex=True)
  plt.scatter(gas.X, gas.Y, marker='+')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  outputfilename=prefix+"Fig01-XY_ini_end"+ext
  plt.savefig(outputfilename)
  plt.clf()

